use crate::action::device::new_serial_number;
use rcgen::{
    BasicConstraints, Certificate, CertificateParams, DistinguishedName, DnType,
    ExtendedKeyUsagePurpose, IsCa, KeyPair, KeyUsagePurpose,
};
use rocket::{
    fairing::{self, Fairing, Info, Kind},
    time::OffsetDateTime,
    Build, Rocket,
};
use rocket_db_pools::{
    sqlx::{self, sqlite::SqliteRow, Row},
    Database,
};
use rustls_pki_types::{CertificateDer, PrivateKeyDer};
use std::{sync::Arc, time::Duration};

#[derive(Database)]
#[database("device")]
pub struct DeviceDb(sqlx::SqlitePool);

pub struct DeviceDbInitializer;

async fn create_device_table(db: &DeviceDb) -> Result<(), sqlx::Error> {
    sqlx::query(
        r#"
            CREATE TABLE IF NOT EXISTS devices (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT NOT NULL,
                certificate TEXT NOT NULL,
                enabled INTEGER NOT NULL,
                UNIQUE (name)
            );
            CREATE TABLE IF NOT EXISTS ca_cert (
                private_key BLOB NOT NULL,
                public_key BLOB NOT NULL,
                certificate BLOB NOT NULL
            );
            CREATE TABLE IF NOT EXISTS rpc_cert (
                private_key BLOB NOT NULL,
                public_key BLOB NOT NULL,
                certificate BLOB NOT NULL
            );
        "#,
    )
    .execute(&db.0)
    .await?;
    Ok(())
}

pub struct CaCertificate(pub Arc<CertificateItem>);
pub struct RpcCertificate(pub Arc<CertificateItem>);

#[derive(Clone)]
pub struct CertificateItem {
    certificate: Arc<Certificate>,
    key_pair: Arc<KeyPair>,
}

impl CertificateItem {
    pub fn certificate(&self) -> &Certificate {
        &self.certificate
    }

    pub fn key_pair(&self) -> &KeyPair {
        &self.key_pair
    }
}

impl From<(Certificate, KeyPair)> for CertificateItem {
    fn from(value: (Certificate, KeyPair)) -> Self {
        CertificateItem {
            certificate: Arc::new(value.0),
            key_pair: Arc::new(value.1),
        }
    }
}

#[derive(Debug)]
#[allow(unused)]
enum InitializeKeyError {
    Database(sqlx::Error),
    Certificate(rcgen::Error),
    Serial(rand::Error),
    Parse(&'static str),
}

impl From<sqlx::Error> for InitializeKeyError {
    fn from(value: sqlx::Error) -> Self {
        InitializeKeyError::Database(value)
    }
}

impl From<rcgen::Error> for InitializeKeyError {
    fn from(value: rcgen::Error) -> Self {
        InitializeKeyError::Certificate(value)
    }
}

impl From<rand::Error> for InitializeKeyError {
    fn from(value: rand::Error) -> Self {
        InitializeKeyError::Serial(value)
    }
}

impl From<&'static str> for InitializeKeyError {
    fn from(value: &'static str) -> Self {
        InitializeKeyError::Parse(value)
    }
}

async fn initialize_ca_key(db: &DeviceDb) -> Result<(Certificate, KeyPair), InitializeKeyError> {
    let ca_key_exists: bool = sqlx::query("SELECT EXISTS (SELECT 1 FROM ca_cert)")
        .try_map(|row: SqliteRow| row.try_get(0))
        .fetch_one(&db.0)
        .await?;

    if !ca_key_exists {
        println!("Generate new key pair for CA.");
        let key_pair = KeyPair::generate()?;

        let (public_key, private_key) = (key_pair.public_key_der(), key_pair.serialize_der());

        let mut certificate_params = CertificateParams::new(vec![])?;

        let mut distinguished_name = DistinguishedName::new();
        distinguished_name.push(DnType::CommonName, "network-scanner-server (CA)");
        certificate_params.distinguished_name = distinguished_name;

        let now = OffsetDateTime::now_utc();
        certificate_params.not_before = now;
        certificate_params.not_after = now + Duration::from_secs(10 * 365 * 24 * 3600);

        certificate_params.serial_number = Some(new_serial_number()?);

        certificate_params.is_ca = IsCa::Ca(BasicConstraints::Constrained(0));

        certificate_params
            .key_usages
            .push(KeyUsagePurpose::KeyCertSign);

        let certificate = certificate_params.self_signed(&key_pair)?;

        sqlx::query(
            r#"
                INSERT INTO ca_cert
                (private_key, public_key, certificate)
                VALUES (?, ?, ?)
            "#,
        )
        .bind(private_key)
        .bind(public_key)
        .bind(certificate.der().as_ref())
        .execute(&db.0)
        .await?;

        Ok((certificate, key_pair))
    } else {
        let key_pair = sqlx::query(
            r#"
                SELECT private_key, public_key
                FROM ca_cert
            "#,
        )
        .map(|row: SqliteRow| {
            let private: Box<[u8]> = row.try_get(0)?;
            let private: PrivateKeyDer = private.as_ref().try_into()?;
            let key_pair = (&private).try_into()?;
            Ok::<KeyPair, InitializeKeyError>(key_pair)
        })
        .fetch_one(&db.0)
        .await??;
        let params = sqlx::query(
            r#"
                SELECT certificate
                FROM ca_cert
            "#,
        )
        .map(|row: SqliteRow| {
            let certificate: Box<[u8]> = row.try_get(0)?;
            let certificate: CertificateDer = certificate.as_ref().into();
            let certificate_params: CertificateParams =
                CertificateParams::from_ca_cert_der(&certificate)?;
            Ok::<CertificateParams, InitializeKeyError>(certificate_params)
        })
        .fetch_one(&db.0)
        .await??;
        let certificate = params.self_signed(&key_pair)?;
        Ok((certificate, key_pair))
    }
}

async fn initialize_rpc_key(
    db: &DeviceDb,
    ca: (&Certificate, &KeyPair),
) -> Result<(Certificate, KeyPair), InitializeKeyError> {
    let rpc_key_exists: bool = sqlx::query("SELECT EXISTS (SELECT 1 FROM rpc_cert)")
        .try_map(|row: SqliteRow| row.try_get(0))
        .fetch_one(&db.0)
        .await?;

    if !rpc_key_exists {
        println!("Generate new key pair for RPC.");
        let key_pair = KeyPair::generate()?;

        let (public_key, private_key) = (key_pair.public_key_der(), key_pair.serialize_der());

        let mut certificate_params =
            CertificateParams::new(vec![String::from("network-scanner-server")])?;

        let mut distinguished_name = DistinguishedName::new();
        distinguished_name.push(DnType::CommonName, "network-scanner-server");
        certificate_params.distinguished_name = distinguished_name;

        let now = OffsetDateTime::now_utc();
        certificate_params.not_before = now;
        certificate_params.not_after = now + Duration::from_secs(10 * 365 * 24 * 3600);

        certificate_params.serial_number = Some(new_serial_number()?);

        certificate_params
            .extended_key_usages
            .push(ExtendedKeyUsagePurpose::ServerAuth);

        let certificate = certificate_params.signed_by(&key_pair, ca.0, ca.1)?;

        sqlx::query(
            r#"
                INSERT INTO rpc_cert
                (private_key, public_key, certificate)
                VALUES (?, ?, ?)
            "#,
        )
        .bind(private_key)
        .bind(public_key)
        .bind(certificate.der().as_ref())
        .execute(&db.0)
        .await?;

        Ok((certificate, key_pair))
    } else {
        let key_pair = sqlx::query(
            r#"
                SELECT private_key, public_key
                FROM rpc_cert
            "#,
        )
        .map(|row: SqliteRow| {
            let private: Box<[u8]> = row.try_get(0)?;
            let private: PrivateKeyDer = private.as_ref().try_into()?;
            let key_pair = (&private).try_into()?;
            Ok::<KeyPair, InitializeKeyError>(key_pair)
        })
        .fetch_one(&db.0)
        .await??;
        let params = sqlx::query(
            r#"
                SELECT certificate
                FROM rpc_cert
            "#,
        )
        .map(|row: SqliteRow| {
            let certificate: Box<[u8]> = row.try_get(0)?;
            let certificate: CertificateDer = certificate.as_ref().into();
            let certificate_params: CertificateParams =
                CertificateParams::from_ca_cert_der(&certificate)?;
            Ok::<CertificateParams, InitializeKeyError>(certificate_params)
        })
        .fetch_one(&db.0)
        .await??;
        let certificate = params.signed_by(&key_pair, ca.0, ca.1)?;
        Ok((certificate, key_pair))
    }
}

#[rocket::async_trait]
impl Fairing for DeviceDbInitializer {
    fn info(&self) -> Info {
        Info {
            name: "Device database initializer",
            kind: Kind::Ignite,
        }
    }

    async fn on_ignite(&self, rocket: Rocket<Build>) -> fairing::Result {
        let db = DeviceDb::fetch(&rocket).expect("Device database is not loaded");
        create_device_table(db)
            .await
            .expect("Can't initialize device database");
        let ca_certificate = initialize_ca_key(db)
            .await
            .expect("Can't initialize CA key");
        let rpc_certificate = initialize_rpc_key(db, (&ca_certificate.0, &ca_certificate.1))
            .await
            .expect("Can't initialize RPC key");
        fairing::Result::Ok(
            rocket
                .manage(CaCertificate(Arc::new(ca_certificate.into())))
                .manage(RpcCertificate(Arc::new(rpc_certificate.into()))),
        )
    }
}
